#include "test.h"
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include "bsl_err.h"
#include "bsl_sal.h"
#include "crypt_errno.h"
#include "crypt_eal_md.h"
#include "crypt_eal_pkey.h"
#include "crypt_eal_init.h"
#include "crypt_eal_rand.h"
#include "crypt_params_key.h"
#include "rng.h"
#include "../../openhitls/crypto/sm4/include/crypt_sm4.h"
#include "crypt_eal_cipher.h"

static void PrintLastError(void) {
    const char *file = NULL;
    uint32_t line = 0;
    BSL_ERR_GetLastErrorFileLine(&file, &line); // 获取错误发生的文件名和行数
    printf("failed at file %s at line %d\r\n", file, (int)line);
}

int32_t sm3_benchmark(void)
{
    int rc = CRYPT_SUCCESS;
    uint8_t in2[] = {
            0x61, 0x62, 0x63, 0x64, 0x61, 0x62, 0x63, 0x64, 0x61, 0x62, 0x63, 0x64, 0x61, 0x62, 0x63, 0x64,
            0x61, 0x62, 0x63, 0x64, 0x61, 0x62, 0x63, 0x64, 0x61, 0x62, 0x63, 0x64, 0x61, 0x62, 0x63, 0x64,
            0x61, 0x62, 0x63, 0x64, 0x61, 0x62, 0x63, 0x64, 0x61, 0x62, 0x63, 0x64, 0x61, 0x62, 0x63, 0x64,
            0x61, 0x62, 0x63, 0x64, 0x61, 0x62, 0x63, 0x64, 0x61, 0x62, 0x63, 0x64, 0x61, 0x62, 0x63, 0x64};

    uint8_t out[32]; // SM3 digest length is 32
    uint32_t outLen = sizeof(out);
    BENCH_TIMES(CRYPT_EAL_Md(CRYPT_MD_SM3, in2, sizeof(in2), out, &outLen), rc, CRYPT_SUCCESS, 10000, "sm3");
    return rc;
}

int32_t sm3_compress_benchmark(void) {
    int rc = CRYPT_SUCCESS;
    uint64_t n = 10000, m = 1;
    uint8_t message[65] = "abc";
    uint32_t digest[8] = {0};
    BENCH_TIMES_VOID(SM3_Compress(digest, message, m), 10000, "sm3 compress");
    return rc;
}

int32_t sm2_keygen_benchmark(void *ctx) {
    int rc = CRYPT_SUCCESS;
    BENCH_TIMES(CRYPT_EAL_PkeyGen(ctx),rc,CRYPT_SUCCESS,10,"sm2 keyGen");
    return rc;
}

static int32_t sm2_sign_inner(void *ctx, int32_t hashId)
{
    uint8_t plainText[32];
    uint8_t signature[256];
    uint32_t signatureLen = sizeof(signature);
    return CRYPT_EAL_PkeySign(ctx, hashId, plainText, sizeof(plainText), signature, &signatureLen);
}

int32_t sm2_sign_benchmark(void *ctx)
{
    int rc;
    BENCH_TIMES(sm2_sign_inner(ctx, CRYPT_MD_SM3), rc, CRYPT_SUCCESS, 10, "sm2 sign");
    return rc;
}

int32_t sm2_verify_benchmark(void *ctx)
{
    int rc;
    uint8_t plainText[32];
    uint8_t signature[256];
    uint32_t signatureLen = sizeof(signature);
    rc = CRYPT_EAL_PkeySign(ctx, CRYPT_MD_SM3, plainText, sizeof(plainText), signature, &signatureLen);
    if (rc != CRYPT_SUCCESS) {
        printf("Failed to sign\r\n");
        return rc;
    }
    BENCH_TIMES(CRYPT_EAL_PkeyVerify(ctx, CRYPT_MD_SM3, plainText, sizeof(plainText), signature, signatureLen), rc,
                CRYPT_SUCCESS, 10, "sm2 verify");
    return rc;
}

static int32_t sm2_enc_inner(void *ctx)
{
    uint8_t plainText[32];
    uint8_t cipherText[256]; // > 32 + 97 + 12
    uint32_t outLen = sizeof(cipherText);
    return CRYPT_EAL_PkeyEncrypt(ctx, plainText, sizeof(plainText), cipherText, &outLen);
}

int32_t sm2_enc_benchmark(void *ctx)
{
    int rc = CRYPT_SUCCESS;
    BENCH_TIMES(sm2_enc_inner(ctx), rc, CRYPT_SUCCESS, 10, "sm2 enc");
    return rc;
}

int32_t sm2_dec_benchmark(void *ctx)
{
    int rc;
    uint8_t plainText[32];
    uint32_t plainTextLen = sizeof(plainText);
    uint8_t cipherText[256]; // > 32 + 97 + 12
    uint32_t outLen = sizeof(cipherText);
    rc = CRYPT_EAL_PkeyEncrypt(ctx, plainText, sizeof(plainText), cipherText, &outLen);
    if (rc != CRYPT_SUCCESS) {
        printf("Failed to encrypt\r\n");
        return rc;
    }
    BENCH_TIMES(CRYPT_EAL_PkeyDecrypt(ctx, cipherText, outLen, plainText, &plainTextLen), rc, CRYPT_SUCCESS, 10, "sm2 dec");
    return rc;
}

void test_benchmarks(void) {
    printf("Starting OpenHitls Crypt Benchmark Test...\r\n");
    BSL_ERR_Init(); // 初始化错误码模块
    // 调用算法API接口之前需要调用BSL_SAL_CallBack_Ctrl函数注册malloc和free函数。该步骤仅需执行一次
    // 如果未注册并且默认能力没有被裁剪,使用默认linux实现
    BSL_SAL_CallBack_Ctrl(BSL_SAL_MEM_MALLOC, (void *)(uintptr_t)malloc);
    BSL_SAL_CallBack_Ctrl(BSL_SAL_MEM_FREE, (void *)(uintptr_t)free);
    if (CRYPT_EAL_Init(CRYPT_EAL_INIT_CPU | CRYPT_EAL_INIT_PROVIDER) != CRYPT_SUCCESS) {
        printf("CRYPT_EAL_Init failed\r\n");
        return;
    }
    int ret;
    uint8_t userId[32] = {0};
    uint8_t key[32] = {0};
    uint8_t msg[32] = {0};
    uint8_t signBuf[100] = {0};
    uint32_t signLen = sizeof(signBuf);
    CRYPT_EAL_PkeyPrv prv = {0};
    CRYPT_EAL_PkeyPub pub = {0};
    CRYPT_EAL_PkeyCtx *ctx = NULL;

    ctx = CRYPT_EAL_PkeyNewCtx(CRYPT_PKEY_SM2);
    if (ctx == NULL) {
        goto EXIT;
    }

    // Set a user ID.
    ret = CRYPT_EAL_PkeyCtrl(ctx, CRYPT_CTRL_SET_SM2_USER_ID, userId, sizeof(userId));
    if (ret != CRYPT_SUCCESS) {
        printf("error code is %x\r\n", ret);
        PrintLastError();
        goto EXIT;
    }

    // --------- 设置熵源回调并注册到Provider ---------
    BSL_Param param[3] = {0};
    BSL_PARAM_InitValue(&param[0], CRYPT_PARAM_RAND_SEED_GETENTROPY,   BSL_PARAM_TYPE_FUNC_PTR, stm32_getEntropy, 0);
    BSL_PARAM_InitValue(&param[1], CRYPT_PARAM_RAND_SEED_CLEANENTROPY, BSL_PARAM_TYPE_FUNC_PTR, stm32_cleanEntropy, 0);
    BSL_PARAM_InitValue(&param[2], CRYPT_PARAM_RAND_SEED_GETNONCE,     BSL_PARAM_TYPE_FUNC_PTR, stm32_getNonce, 0);
    BSL_PARAM_InitValue(&param[3], CRYPT_PARAM_RAND_SEED_CLEANNONCE,   BSL_PARAM_TYPE_FUNC_PTR, stm32_cleanNonce, 0);

    ret = CRYPT_EAL_ProviderRandInitCtx(NULL, CRYPT_RAND_SHA256, "provider=default", NULL, 0, param);

    CRYPT_EAL_SetRandCallBack((CRYPT_EAL_RandFunc)myRand);
    if (ret != CRYPT_SUCCESS) {
        printf("error code is %x\r\n", ret);
        PrintLastError();
        goto EXIT;
    }
    // Generate a key pair.
    ret = CRYPT_EAL_PkeyGen(ctx);
    if (ret != CRYPT_SUCCESS) {
        printf("error code is %x\r\n", ret);
        PrintLastError();
        goto EXIT;
    }

    ret = CRYPT_EAL_PkeySign(ctx, CRYPT_MD_SM3, msg, sizeof(msg), signBuf, &signLen);
    if (ret != CRYPT_SUCCESS) {
        printf("error code is %x\r\n", ret);
        PrintLastError();
        goto EXIT;
    }

    // Verify the signature.
    ret = CRYPT_EAL_PkeyVerify(ctx, CRYPT_MD_SM3, msg, sizeof(msg), signBuf, signLen);
    if (ret != CRYPT_SUCCESS) {
        printf("error code is %x\r\n", ret);
        PrintLastError();
        goto EXIT;
    }
    printf("%-25s, %15s, %20s, %20s\r\n", "algorithm operation", "time elapsed(s)", "run times", "ops");
    sm3_benchmark();
    sm3_compress_benchmark();

    sm2_keygen_benchmark(ctx);

    sm2_sign_benchmark(ctx);

    sm2_verify_benchmark(ctx);

    sm2_enc_benchmark(ctx);

    sm2_dec_benchmark(ctx);

    printf("pass \r\n");
    EXIT:
    // Release the context memory.
    CRYPT_EAL_PkeyFreeCtx(ctx);
    CRYPT_EAL_RandDeinit();
    BSL_ERR_DeInit();
}
